import pandas as pd
from mlxtend.preprocessing import TransactionEncoder
from mlxtend.frequent_patterns import apriori, association_rules
import matplotlib.pyplot as plt
import networkx as nx

# 加载数据
df = pd.read_excel('餐厅数据.xlsx')

# 转换菜品数据格式
transactions = df['菜品'].str.split(',').to_list()

te = TransactionEncoder()
te_ary = te.fit(transactions).transform(transactions)
df_encoded = pd.DataFrame(te_ary, columns=te.columns_)

# 使用apriori进行关联规则挖掘
frequent_itemsets = apriori(df_encoded, min_support=0.1, use_colnames=True)
frequent_itemsets.sort_values(by='support', ascending=False, inplace=True)

# 生成关联规则
rules = association_rules(frequent_itemsets, metric='confidence', min_threshold=0.1)
# 筛选有效规则
effective = rules[
    (rules['lift'] > 1) & (rules['confidence'] > 0.1)
].sort_values(by=['lift', 'confidence'], ascending=False)

# 查看筛选后的规则
print(effective[['antecedents', 'consequents', 'support', 'confidence', 'lift']])

# 设置字体
plt.rcParams['font.sans-serif'] = ['Simhei']
plt.rcParams['axes.unicode_minus'] = False

# 关联关系可视化
G = nx.DiGraph()
for _, row in effective.iterrows():
    G.add_edge(','.join(list(row['antecedents'])),
               ','.join(list(row['consequents'])),
               weight=row['lift'])

plt.figure(figsize=(12, 8))
pos = nx.spring_layout(G, k=0.5)  # k值越小，节点间距越近

# 绘制边，通过边的宽度来表示lift值的不同
edges = G.edges()
weights = [G[u][v]['weight'] for u, v in edges]

nx.draw(G, pos,
        with_labels=True,
        edge_color='blue',
        width=[w * 2 for w in weights],  # 根据lift值调整边的宽度
        node_color='red',
        node_size=500,
        font_size=10,
        arrowsize=20)

plt.title("关联规则网络图")
plt.show()